{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset    # Dataset读取单个样本的抽象类，包括读取样本的标签、样本的索引、样本的特征\n",
    "from torch.utils.data import DataLoader   # 数据加载器，组合数据集和采样器，并在数据集上提供单进程或多进程迭代器\n",
    "from torchvision import datasets    # torchvision是PyTorch中专门用于图像处理的库,datasets是其中的数据集模块\n",
    "from torchvision.transforms import ToTensor\n",
    "\n",
    "from torch.nn import functional as F\n",
    "from torch import nn, optim\n",
    "\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_curve(data):\n",
    "    \"\"\"\n",
    "    下降曲线的绘制\n",
    "    :param data:\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    # fig = plt.figure()\n",
    "    plt.plot(range(len(data)), data, color='blue')  #range(len(data))用于生成一个从0到len(data)的整数序列表示横坐标，data表示纵坐标，color表示颜色\n",
    "    plt.legend(['value'], loc='upper right')#图例，loc表示位置\n",
    "    plt.xlabel('step')#横坐标\n",
    "    plt.ylabel('value')#纵坐标\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def plot_image(img, label, name):\n",
    "    # fig = plt.figure()\n",
    "    for i in range(6):\n",
    "        plt.subplot(2, 3, i + 1)     # 2行3列,i+1表示第几个子图\n",
    "        plt.tight_layout()  # 自动调整子图参数，使之填充整个图像区域\n",
    "        plt.imshow(img[i][0], cmap='gray', interpolation='none')  # img0维度存储的是图片，i是第i张图片。interpolation表示插值方法\n",
    "        plt.title(\"{}: {}\".format(name, label[i].item())) # 显示标题，label[i].item()表示取出label中的第i个元素，item()将tensor转换为标量，用于画图\n",
    "        plt.xticks([])\n",
    "        plt.yticks([])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_dataset():\n",
    "    #从torchversion加载自带数据集\n",
    "    training_data = datasets.MNIST(\n",
    "    root=\"../data\",    # 数据集存放路径\n",
    "    train=True,    # 是否为训练集\n",
    "    download=True,    # 是否下载\n",
    "    transform=ToTensor(),    # 数据转换方式\n",
    "    )\n",
    "    train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)    # DataLoader是一个迭代器，每次返回一个batch的数据\n",
    "    test_data = datasets.MNIST(\n",
    "        root=\"../data\",\n",
    "        train=False,\n",
    "        download=True,\n",
    "        transform=ToTensor()\n",
    "    )\n",
    "\n",
    "    test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)  \n",
    "    return train_dataloader, test_dataloader\n",
    "#打印datasets支持的数据集\n",
    "print(datasets.__all__)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        # xw + b\n",
    "        self.fc1 = nn.Linear(28 * 28, 256)\n",
    "        self.fc2 = nn.Linear(256, 64)\n",
    "        self.fc3 = nn.Linear(64, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # x: [b, 1, 28, 28]\n",
    "        # h1 = relu(xw1 + b)\n",
    "        input=x #保存输入\n",
    "        x = F.relu(self.fc1(x))\n",
    "        # h2 = relu(h1w2 + b2)\n",
    "        x = F.relu(self.fc2(x))\n",
    "        # h3 = h2w3 + b3\n",
    "        x = self.fc3(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch_num, train_loader, net, optimizer):\n",
    "    train_loss = []\n",
    "    for epoch in range(epoch_num):\n",
    "        for batch_idx, (x, y) in enumerate(train_loader):\n",
    "            # x: [b, 1, 28, 28], y: [512]\n",
    "            # [b, 1, 28, 28] => [b, feature]\n",
    "            x = x.view(x.size(0), 28 * 28)\n",
    "            # => [b, 10]\n",
    "            out = net(x)\n",
    "            # y_onehot = one_hot(y) # [b] => [b, 10]\n",
    "            y_onehot = F.one_hot(y, 10).float()\n",
    "            # loss = mse(out, y_onehot)\n",
    "            loss = F.mse_loss(out, y_onehot)\n",
    "            # print(loss)\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            # w' = w - lr * grad\n",
    "            optimizer.step()\n",
    "\n",
    "            train_loss.append(loss.item())\n",
    "\n",
    "            if batch_idx % 100 == 0:\n",
    "                print('epoch {}, batch_idx {}, loss {}'.format(epoch, batch_idx, loss.item()))\n",
    "\n",
    "    return train_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "train_loader, test_loader = load_dataset()\n",
    "x, y = next(iter(train_loader))\n",
    "print(x.shape, y.shape, x.min(), x.max())\n",
    "plot_image(x, y, \"image sample\")\n",
    "\n",
    "net = Net() #调用Net类\n",
    "# [w1, b1, w2, b2, w3, b3]\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)  #定义优化器，参数说明：net.parameters()表示优化的参数。\n",
    "\n",
    "epoch_num = 1   #训练次数1\n",
    "train_loss = train(epoch_num, train_loader, net, optimizer) #训练返回loss\n",
    "\n",
    "plot_curve(train_loss)  #绘制下降曲线\n",
    "# 得到参数[w1, b1, w2, b2, w3, b3]\n",
    "\n",
    "total_correct = 0   #初始化一个变量，用于记录正确的个数\n",
    "for x, y in test_loader:\n",
    "    x = x.view(x.size(0), 28 * 28)\n",
    "    out = net(x)\n",
    "    # out: [b, 10] => pred: [b]\n",
    "    pred = out.argmax(dim=1)    #返回最大值的索引\n",
    "    correct = pred.eq(y).sum().float().item()   #计算正确的个数，eq()比较两个tensor是否相等，sum()求和，float()转换为浮点数，item()返回元素值\n",
    "    total_correct += correct    #累加正确的个数\n",
    "\n",
    "total_num = len(test_loader.dataset)    #测试集的总数\n",
    "acc = total_correct / total_num #计算准确率\n",
    "print('test acc:', acc)\n",
    "\n",
    "x, y = next(iter(test_loader))  #取出一个batch的数据，iter()返回一个迭代器，next()返回迭代器的下一个项目？？\n",
    "out = net(x.view(x.size(0), 28 * 28))   #预测，此时out是一个tensor，存储的是一组预测值，每个值是一个10维的向量\n",
    "pred = out.argmax(dim=1)    #取出预测值中最大的索引，即预测的类别\n",
    "plot_image(x, pred, 'test') #绘制图像"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dl2024",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
